import numpy as np
from numpy.typing import NDArray
from scipy.spatial.transform import Rotation as scipy_R
from typing import Union, List
from shapely.geometry import LineString, Polygon, Point

DEFAULT_ALL_CLOSE_PRECISION = 1.e-4

def euler_to_rotation_matrix(euler_xyz: Union[NDArray, list], degree=False) -> NDArray:
    return scipy_R.from_euler('xyz', euler_xyz, degrees=degree).as_matrix()

def rotation_matrix_to_euler(rotation_matrix: NDArray) -> NDArray:
    return scipy_R.from_matrix(rotation_matrix).as_euler('xyz')

def quaternion_to_rotation_matrix(quat_xyzw) -> NDArray:
    return scipy_R.from_quat(quat_xyzw).as_matrix()

def quaternion_to_euler(quat_xyzw) -> NDArray:
    return scipy_R.from_quat(quat_xyzw).as_euler('xyz')

def check_is_3d_rotation_matrix(matrix: NDArray, precision=DEFAULT_ALL_CLOSE_PRECISION) -> bool:
    if matrix.shape != (3, 3):
        return False
    return np.allclose(matrix.dot(matrix.T), np.identity(3), atol=precision) and \
        np.allclose(np.linalg.det(matrix), 1, atol=precision)

def check_is_3d_transform_matrix(matrix: NDArray, precision=DEFAULT_ALL_CLOSE_PRECISION) -> bool:
    if matrix.shape != (4, 4):
        return False
    return check_is_3d_rotation_matrix(matrix[:3, :3], precision) and \
        np.allclose(matrix[3, :], np.array([0, 0, 0, 1]), atol=precision)

def check_is_camera_intrinsic(matrix: NDArray, precision=DEFAULT_ALL_CLOSE_PRECISION) -> bool:
    if matrix.shape != (3, 3):
        return False
    return np.allclose(matrix[2, :], np.array([0, 0, 1]), atol=precision)

def transform_3d_cloud(cloud: NDArray, translation: NDArray,
                       rotation_matrix: NDArray, inverse: bool=False) -> NDArray:
    """ transform 3D point cloud while keeping cloud's shape
    """
    assert translation.size == 3
    assert check_is_3d_rotation_matrix(rotation_matrix)
    assert cloud.shape[1] > 2
    cloud_ = cloud.copy()
    translation = translation.reshape((1, 3))
    if inverse:
        tmp_xyz = cloud_[:, :3] - translation
        cloud_[:, :3] = rotation_matrix.T.dot(tmp_xyz.T).T
    else:
        cloud_[:, :3] = rotation_matrix.dot(cloud_[:, :3].T).T + translation
    return cloud_

def transform_3d_cloud_to_image(cloud: NDArray, cloud_to_camera: NDArray,
                                camera_intrinsic: NDArray) -> NDArray:
    """ transform 3D point cloud to image pixels while keeping cloud's shape
    """
    assert check_is_3d_transform_matrix(cloud_to_camera)
    assert check_is_camera_intrinsic(camera_intrinsic)
    assert cloud.shape[1] > 2 and len(cloud.shape) == 2
    cloud = transform_3d_cloud(cloud, cloud_to_camera[:3, 3], cloud_to_camera[:3, :3])
    homo_xyz = cloud[:, :3] / cloud[:, 2].reshape(-1, 1)
    cloud[:, :2] = camera_intrinsic.dot(homo_xyz.T).T[:, :2]
    return cloud

def bounding_box_to_points(position: list, size: list, rotation: list) -> NDArray:
    """ calculate 3D bounding box's points
    :param position: position of box
    :param size: size of box (x, y, z)
    :param rotation: heading of box in euler angles
    """
    assert len(position) == len(size) == len(rotation) == 3
    pos_arr = np.array(position)
    size_arr = np.array(size)
    pts_prefix = np.array([[-1, -1, -1], [-1, 1, -1], [1, 1, -1], [1, -1, -1],
                           [-1, -1, 1], [-1, 1, 1], [1, 1, 1], [1, -1, 1]]) / 2
    return euler_to_rotation_matrix(rotation).dot((pts_prefix * size_arr).T).T + pos_arr

def point_in_polygon_roi(point: NDArray, roi_pts: NDArray) -> bool:
    assert point.size <= 3 and roi_pts.shape[1] <= 3
    roi = Polygon(roi_pts)
    return roi.contains(Point(*point))

def line_in_polygon_roi(line_pts: NDArray, roi_pts: NDArray) -> List[NDArray]:
    assert line_pts.shape[1] <= 3 and roi_pts.shape[1] <= 3
    line = LineString(line_pts)
    roi = Polygon(roi_pts)
    inter = line.intersection(roi)
    if isinstance(inter, LineString):
        return [np.array(inter.coords)]
    else:
        return [np.array(geo.coords) for geo in inter.geoms]

def polyline_in_polygon_roi(poly_pts: NDArray, roi_pts: NDArray) -> List[NDArray]:
    assert poly_pts.shape[1] <= 3 and roi_pts.shape[1] <= 3
    poly = Polygon(poly_pts).boundary
    roi = Polygon(roi_pts)
    inter = poly.intersection(roi)
    if isinstance(inter, LineString):
        return [np.array(inter.coords)]
    else:
        return [np.array(geo.coords) for geo in inter.geoms]
